import tf from '../tf'
import { RRDB } from './RRDB'

/**
 * Make a layer with multiple RRDB blocks
 */
function makeLayer(numBlock: number, numFeat: number, numGrowCh: number) {
  const blocks = [] as RRDB[]
  for (let i = 0; i < numBlock; i++) {
    blocks.push(new RRDB(numFeat, numGrowCh))
  }
  return blocks
}

interface RRDBNetParams {
  numInCh: number
  numOutCh: number
  scale?: number
  numFeat?: number
  numBlock?: number
  numGrowCh?: number
}
const defaultPamams: RRDBNetParams = {
  numInCh: 3,
  numOutCh: 3,
  scale: 4,
  numFeat: 64,
  numBlock: 23,
  numGrowCh: 32,
}
/**
 * RRDBNet model
 */
export class RRDBNet extends tf.layers.Layer {
  scale: number
  numInCh: number
  numOutCh: number
  numFeat: number
  convFirst: tf.layers.Layer
  body: RRDB[]
  convBody: tf.layers.Layer
  convUp1: tf.layers.Layer
  convUp2: tf.layers.Layer
  convHr: tf.layers.Layer
  convLast: tf.layers.Layer
  lrelu: tf.layers.Layer
  constructor(opt: RRDBNetParams) {
    super({})
    const { numInCh, numOutCh, scale, numFeat, numBlock, numGrowCh } = {
      ...defaultPamams,
      ...opt,
    } as Required<RRDBNetParams>
    this.scale = scale
    this.numInCh = numInCh
    this.numFeat = numFeat
    if (scale === 2) {
      this.numInCh *= 4
    } else if (scale === 1) {
      this.numInCh *= 16
    }
    this.numOutCh = numOutCh
    this.convFirst = tf.layers.conv2d({
      filters: numFeat,
      kernelSize: 3,
      padding: 'same',
    })
    this.body = makeLayer(numBlock, numFeat, numGrowCh)
    this.convBody = tf.layers.conv2d({
      filters: numFeat,
      kernelSize: 3,
      padding: 'same',
    })
    this.convUp1 = tf.layers.conv2d({
      filters: numFeat,
      kernelSize: 3,
      padding: 'same',
    })
    this.convUp2 = tf.layers.conv2d({
      filters: numFeat,
      kernelSize: 3,
      padding: 'same',
    })
    this.convHr = tf.layers.conv2d({
      filters: numFeat,
      kernelSize: 3,
      padding: 'same',
    })
    this.convLast = tf.layers.conv2d({
      filters: numOutCh,
      kernelSize: 3,
      padding: 'same',
    })
    this.lrelu = tf.layers.leakyReLU({ alpha: 0.2 })
  }

  call(x: any) {
    let feat
    if (this.scale === 2) {
      feat = this.pixelUnshuffle(x, 2)
    } else if (this.scale === 1) {
      feat = this.pixelUnshuffle(x, 4)
    } else {
      feat = x
    }
    feat = this.convFirst.apply(feat) as any
    let bodyFeat = feat
    for (const block of this.body) {
      bodyFeat = block.apply(bodyFeat)
    }

    bodyFeat = this.convBody.apply(bodyFeat)
    feat = feat.add(bodyFeat)
    feat = this.lrelu.apply(
      this.convUp1.apply(
        tf.image.resizeNearestNeighbor(feat, [
          feat.shape[1] * 2,
          feat.shape[2] * 2,
        ])
      )
    )
    feat = this.lrelu.apply(
      this.convUp2.apply(
        tf.image.resizeNearestNeighbor(feat as tf.Tensor3D, [
          // @ts-ignore
          feat.shape[1] * 2,
          // @ts-ignore
          feat.shape[2] * 2,
        ])
      )
    )
    const out = this.convLast.apply(this.lrelu.apply(this.convHr.apply(feat)))
    return out as any
  }
  override computeOutputShape(inputShape: tf.Shape): tf.Shape | tf.Shape[] {
    if (inputShape == null) {
      throw new Error('inputShape is null')
    }
    // if (inputShape[1] == null || inputShape[2] == null) {
    //   throw new Error('inputShape[1] or inputShape[2] is null')
    // }
    return [
      inputShape[0],
      // @ts-ignore
      inputShape[1] * this.scale || null,
      // @ts-ignore
      inputShape[2] * this.scale || null,
      inputShape[3],
    ]
  }

  pixelUnshuffle(x: tf.Tensor4D, scale: number) {
    if (Array.isArray(x)) {
      x = x[0]
    }
    const batchSize = x.shape[0]
    const inHeight = x.shape[1]
    const inWidth = x.shape[2]
    const inChannels = x.shape[3]
    const outChannels = inChannels * scale * scale
    const outHeight = inHeight / scale
    const outWidth = inWidth / scale
    const xReshaped = x.reshape([
      batchSize,
      outHeight,
      scale,
      outWidth,
      scale,
      inChannels,
    ])
    const xTransposed = xReshaped.transpose([0, 1, 3, 2, 4, 5])
    const xUnshuffled = xTransposed.reshape([
      batchSize,
      outHeight,
      outWidth,
      outChannels,
    ])

    return xUnshuffled
  }
  build(inputShape: tf.Shape): void {
    const trainableWeights: tf.LayerVariable[] = []
    if (this.scale === 2) {
      inputShape[2]! *= 4
    } else if (this.scale === 1) {
      inputShape[2]! *= 16
    }
    inputShape[2] = inputShape[2] || null
    this.convFirst.build(inputShape)
    trainableWeights.push(...this.convFirst.trainableWeights)
    for (const block of this.body) {
      block.build(inputShape)
      trainableWeights.push(...block.trainableWeights)
    }
    this.convBody.build([
      inputShape[0],
      inputShape[1],
      inputShape[2],
      this.numFeat,
    ])
    trainableWeights.push(...this.convBody.trainableWeights)
    this.convUp1.build([null, null, null, this.numFeat])
    trainableWeights.push(...this.convUp1.trainableWeights)
    this.convUp2.build([null, null, null, this.numFeat])
    trainableWeights.push(...this.convUp2.trainableWeights)
    this.convHr.build([null, null, null, this.numFeat])
    trainableWeights.push(...this.convHr.trainableWeights)
    this.convLast.build([null, null, null, this.numFeat])
    trainableWeights.push(...this.convLast.trainableWeights)
    this._trainableWeights = trainableWeights
    this.built = true
  }
  override getWeights(trainableOnly?: boolean): tf.Tensor[] {
    const weights: tf.Tensor[] = []
    weights.push(...this.convFirst.getWeights(trainableOnly))
    for (const block of this.body) {
      weights.push(...block.getWeights(trainableOnly))
    }
    weights.push(...this.convBody.getWeights(trainableOnly))
    weights.push(...this.convUp1.getWeights(trainableOnly))
    weights.push(...this.convUp2.getWeights(trainableOnly))
    weights.push(...this.convHr.getWeights(trainableOnly))
    weights.push(...this.convLast.getWeights(trainableOnly))
    return weights
  }
  override setWeights(weights: tf.Tensor[]): void {
    weights = [...weights]
    this.convFirst.setWeights(weights.splice(0, 2))
    // const blockWeight = weights.splice(2, 2 + this.body.length * 2)
    for (const block of this.body) {
      block.setWeights(weights.splice(0, 2))
    }
    this.convBody.setWeights(weights.splice(0, 2))
    this.convUp1.setWeights(weights.splice(0, 2))
    this.convUp2.setWeights(weights.splice(0, 2))
    this.convHr.setWeights(weights.splice(0, 2))
    this.convLast.setWeights(weights.splice(0, 2))
  }
  static get className() {
    return 'RRDBNet'
  }
}
tf.serialization.registerClass(RRDBNet)
// Example usage:
// const model = new RRDBNet({
//   numInCh: 3,
//   numOutCh: 3,
//   scale: 4,
//   numBlock: 6,
// })
// const rRDBNet = new RRDBNet({ numInCh: 3, numOutCh: 3, scale: 4, numBlock: 6 })
// const _input = tf.input({ shape: [64, 64, 3] })
// const model = tf.model({
//   inputs: _input,
//   outputs: rRDBNet.apply(_input) as any,
// })
// console.log(model.getWeights())
// const output = model.predict(input) as any
// console.log(output.shape)
